--- Test case generation.
module Test.QuickCheckGen where

-- ------------------------------------------------------------------------
-- imports

import System.Random
  ( Random
  , StdGen
  -- , randomR
  , split
  , newStdGen
  )


import Prelude.Math (round, log)

-- ------------------------------------------------------------------------
--- Generator type

data Gen a = MkGen{ !unGen :: StdGen -> Int -> a }

instance Functor Gen where
  fmap f (MkGen h) =
    MkGen (\r\n -> f (h r n))



instance Monad Gen where
  (<*>) = ap

  pure x =
    MkGen (\_ \_ -> x)

  MkGen m >>= k =
    MkGen (\r \n ->
      let (r1,r2)  = split r
          MkGen m' = k (m r1 n)
       in m' r2 n
    )

-- ------------------------------------------------------------------------
-- ** Primitive generator combinators

-- | Modifies a generator using an integer seed.
variant :: Integral n => n -> Gen a -> Gen a
variant k0 (MkGen m) = MkGen (\r \n -> m (var k0 r) n)
 where
  -- var :: Integral n => n 
  var :: (Random.RandomGen α, Integral n) => n -> α->α
  var k = (if k == k' then id  else var k')
        . (if even k  then fst else snd)
        . split
   where
    k' = k `div` 2

--- Used to construct generators that depend on the size parameter.
sized :: (Int -> Gen a) -> Gen a
sized f = MkGen (\r \n -> let MkGen m = f n in m r n)

{-- 
    Overrides the size parameter. Returns a generator which uses
    the given size instead of the runtime-size parameter.
-}
resize :: Int -> Gen a -> Gen a
resize n (MkGen m) = MkGen (\r \_ -> m r n)

--- Generates a random element in the given inclusive range.
choose :: Random a => (a,a) -> Gen a
choose rng = MkGen (\r \_ -> let (x,_) = randomR rng r in x)

--- Promotes a monadic generator to a generator of monadic values.
promote :: Monad m => m (Gen a) -> Gen (m a)
promote m = MkGen (\r \n -> liftM (\(MkGen m') -> m' r n) m)

--- Generates some example values.
sample' :: Gen a -> IO [a]
sample' (MkGen m) =
  do rnd0 <- newStdGen
     let rnds rnd = rnd1 : rnds rnd2 where (rnd1,rnd2) = split rnd
     return [(m r n) | (r,n) <- rnds rnd0 `zip` [0,2..20] ]

--- Generates some example values and prints them to 'stdout'.
sample :: Show a => Gen a -> IO ()
sample g =
  do cases <- sample' g
     sequence_ (map println cases)

-- ------------------------------------------------------------------------
-- ** Common generator combinators

--- Generates a value that satisfies a predicate.
suchThat :: Gen a -> (a -> Bool) -> Gen a
gen `suchThat` p =
  do mx <- gen `suchThatMaybe` p
     case mx of
       Just x  -> return x
       Nothing -> sized (\n -> resize (n+1) (gen `suchThat` p))

--- Tries to generate a value that satisfies a predicate.
suchThatMaybe :: Gen a -> (a -> Bool) -> Gen (Maybe a)
gen `suchThatMaybe` p = sized (try 0 . max 1)
 where
  try _ 0 = return Nothing
  try k n = do x <- resize (2*k+n) gen
               if p x then return (Just x) else try (k+1) (n-1)

--- Randomly uses one of the given generators. The input list must be non-empty.
oneof :: [Gen a] -> Gen a
oneof [] = error "QuickCheck.oneof used with empty list"
oneof gs = choose (0,length gs - 1) >>= (gs !!)

{-- Chooses one of the given generators, with a weighted random distribution.
    The input list must be non-empty.
    -}
frequency :: [(Int, Gen a)] -> Gen a
frequency [] = error "QuickCheck.frequency used with empty list"
frequency xs0 = choose (1, tot) >>= (`pick` xs0)
 where
  tot = sum (map fst xs0)

  pick n ((k,x):xs)
    | n <= k    = x
    | otherwise = pick (n-k) xs
  pick _ _  = error "QuickCheck.pick used with empty list"

--- Generates one of the given values. The input list must be non-empty.
elements :: [a] -> Gen a
elements [] = error "QuickCheck.elements used with empty list"
elements xs = (xs !!) `fmap` choose (0, length xs - 1)

{-- 
    Takes a list of elements of increasing size, and chooses
    among an initial segment of the list. The size of this initial
    segment increases with the size parameter.
    The input list must be non-empty. -}
growingElements :: [a] -> Gen a
growingElements [] = error "QuickCheck.growingElements used with empty list"
growingElements xs = sized ( \n -> elements (take (1 `max` size n) xs)  )
  where
   k      = length xs
   mx     = 100
   log'   = Long.int . round . log . Int.float
   size n = (log' n + 1) * k `div` log' mx

{- WAS:
growingElements xs = sized $ \n -> elements (take (1 `max` (n * k `div` 100)) xs)
 where
  k = length xs
-}

--- Generates a list of random length. 
--- The maximum length depends on the size parameter.
listOf :: Gen a -> Gen [a]
listOf gen = sized ( \n ->
          do k <- choose (0,n)
             vectorOf k gen
        )
        
--- Generates a non-empty list of random length. 
--- The maximum length depends on the size parameter.
listOf1 :: Gen a -> Gen [a]
listOf1 gen = sized ( \n ->
          do k <- choose (1,1 `max` n)
             vectorOf k gen
    )

--- Generates a list of the given length.
vectorOf :: Int -> Gen a -> Gen [a]
vectorOf k gen = sequence [ gen | _ <- [1..k] ]
